from paddlenlp.transformers import BertModel, BertForTokenClassification
from paddle.nn import Layer, Linear, Softmax
import paddle


class PoetryBertModel(Layer):
    """
    基于BERT预训练模型的诗歌生成模型
    """



    def __init__(self, pretrained_bert_model: str, input_length: int):
        super(PoetryBertModel, self).__init__()
        bert_model = BertModel.from_pretrained(pretrained_bert_model)
        self.vocab_size, self.hidden_size = bert_model.embeddings.word_embeddings.parameters()[0].shape
        # 进行命名实体识别的方法。就是BertForTokenClassfication类
        # 这里必须要写num_labels参数，不然回去自动下载报错
        self.bert_for_class = BertForTokenClassification.from_pretrained(pretrained_bert_model,
                                                                         num_labels=self.vocab_size)
        # 生成下三角矩阵，用来mask句子后边的信息
        self.sequence_length = input_length
        # lower_triangle_mask为input_length * input_length的下三角矩阵（包含主对角线），该掩码作为注意力掩码的一部分（在forward的
        # 处理中为0的部分会被处理成无穷小量，以方便在计算注意力权重的时候保证被掩盖的部分权重约等于0）。而之所以写为下三角矩阵的形式，与
        # transformer的多头注意力计算的机制有关，细节可以了解相关论文获悉。
        '''
        1.在计算注意力权重时,对于位置i和位置j,如果j > i,则表示位置j在位置i之后,我们希望位置i不能关注到位置j的信息,因为这会造成信息泄露。
        2.为了实现这一点,我们会引入一个下三角掩码矩阵,其中下三角元素为0,上三角元素为负无穷。 
        3.将这个下三角掩码矩阵与注意力权重矩阵相加,就可以将上三角部分的注意力权重设为负无穷,从而达到只允许当前位置关注到之前位置的目的。
        4.这种做法可以确保transformer模型在生成序列时,只能关注到当前位置之前的信息,而不会"窥视"未来的信息,从而保证了模型的自回归性质。
        5.因此,下三角矩阵在transformer的多头注意力机制中扮演着重要的角色,是实现自回归性的关键技术之一。
            这也是transformer模型相比于传统RNN在序列生成任务上取得成功的一个重要原因。
        '''
        self.lower_triangle_mask = paddle.tril(paddle.tensor.full((input_length, input_length),
                                                                  1,
                                                                  'float32'))
        """
        tril    返回输入矩阵 input 的下三角部分，其余部分被设为0。 矩形的下三角部分被定义为对角线上和下方的元素。
        full    创建形状大小为 shape 并且数据类型为 dtype 的 Tensor，其中元素值均为 fill_value 
        cast    将 x 的数据类型转换为 dtype 并输出。支持输出和输入的数据类型相同。
        matmul  计算两个 Tensor 的乘积，遵循完整的广播规则，关于广播规则
        """

    def forward(self, token, token_type, input_mask, input_length=None):
        # print("token-shape",token.shape)        # type-shape [1, 128]
        # print("token_type-shape", token_type.shape)   # [1, 128]
        '''
            实际序列长度会比原始序列多一是因为序列以[CLS]开头会填充
        '''
        # print("token_type", token_type)
        # 计算attention mask
        # print("input_mask.shape:",input_mask.shape)         # input_mask.shape: [1, 128]
        mask_left = paddle.reshape(input_mask, input_mask.shape + [1])
        mask_right = paddle.reshape(input_mask, [input_mask.shape[0], 1, input_mask.shape[1]])
        # 输入句子中有效的位置
        mask_left = paddle.cast(mask_left, 'float32')
        mask_right = paddle.cast(mask_right, 'float32')
        # print("mask_left.shape", mask_left.shape)           # mask_left [1, 128, 1]
        # print("mask_right.shape", mask_right.shape)         # mask_right [1, 1, 128]
        # print("mask_left", mask_left)
        # print("mask_right", mask_right)
        attention_mask = paddle.matmul(mask_left, mask_right)
        # print("attention_mask1.shape", attention_mask.shape)       # attention_mask [1, 128, 128]
        '''
            可以缩小input_length和输入的诗词长度，查看完整的矩阵，理解掩码和下三角矩阵
        '''
        # print("attention_mask", attention_mask)
        # 注意力机制计算中有效的位置
        if input_length is not None:
            # 之所以要再计算一次，是因为用于推理预测时，可能输入的长度不为实例化时设置的长度。这里的模型在训练时假设输入的
            # 长度是被填充成一致的——这一步不是必须的，但是处理成一致长度比较方便处理（对应地，增加了显存的用度）。
            lower_triangle_mask = paddle.tril(paddle.tensor.full((input_length, input_length), 1, 'float32'))
            # print("lower_triangle_mask.shape", lower_triangle_mask.shape)  # [128, 128]
            # print("lower_triangle_mask", lower_triangle_mask)
        else:
            lower_triangle_mask = self.lower_triangle_mask
        attention_mask = attention_mask * lower_triangle_mask
        # print("attention_mask-2", attention_mask)
        # 无效的位置设为极小值
        # print("attention_mask2.shape",attention_mask.shape)        # [1, 128, 128]
        attention_mask = (1 - paddle.unsqueeze(attention_mask, axis=[1])) * -1e10
        # print("attention_mask3.shape",attention_mask.shape)        # [1, 1, 128, 128]
        attention_mask = paddle.cast(attention_mask, self.bert_for_class.parameters()[0].dtype)
        # print("self.bert_for_class.parameters()[0].dtype:",self.bert_for_class.parameters()[0].dtype)  # paddle.float32
        output_logits = self.bert_for_class(token, token_type_ids=token_type, attention_mask=attention_mask)
        return output_logits


if __name__ == '__main__':
    from PoemData import PoemData
    import numpy as np
    from paddlenlp.transformers import BertTokenizer

    bert_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
    PoemData = PoemData(
        poems=[
            "松下中峰路，疎阴偃盖清。远云收海雨，数鴈过秋城。抚事心犹壮，卑栖气未平。近来唯此乐，高韵寄闲情。",
            "还源去，欲求般若易。但息是非心，自然成大智。",
            "一山夐出众山中，日照烟浮紫翠浓。山不争高人共仰，几多丘垤俨相从。",
            "还源去，欲求般若易。但息是非心，自然成大智。",
            "客从南方来，酌我一瓯茗。我酌初不啜，彊啜且复醒。既凿浑沌氏，遂出华胥境。操戈逐儒生，举觞还酩酊。",
            "测试长度为十"
        ],
        tokenizer=bert_tokenizer,
        max_len=10)
    # print(PoemData.__len__())
    input_token1, input_token_type1, input_pad_mask1, label_token1, input_pad_mask1 = PoemData.__getitem__(5)
    # print("input_token1-row:", input_token1)
    # print("input_token1decode:", bert_tokenizer.decode(input_token1))
    # print("label_token1decode:", bert_tokenizer.decode(label_token1))
    known_ids = paddle.to_tensor([input_token1], dtype='int64')
    type_token = paddle.to_tensor([input_token_type1], dtype='int64')
    mask = paddle.to_tensor([input_pad_mask1], dtype='int64')
    net = PoetryBertModel('bert-base-chinese', 10)
    logits = net.forward(known_ids, type_token, mask, 10)
    # print("logits:", logits)    # Tensor(shape=[1, 128, 21128]
    words_prob = logits[0, -1, :].numpy()
    # print("words_prob:", words_prob)  # shape (21128,)
    # 依概率倒序排列后，选取前topk个词
    topk = 2
    words_to_be_choosen = words_prob.argsort()[::-1][:topk]
    # print("words_to_be_choosen:", words_to_be_choosen)      # [14435  2742]
    probs_to_be_choosen = words_prob[words_to_be_choosen]
    # print("probs_to_be_choosen-1:", probs_to_be_choosen)    # [2.1776009 1.7892948]
    # 归一化
    probs_to_be_choosen = probs_to_be_choosen / sum(probs_to_be_choosen)
    # print("probs_to_be_choosen-2:", probs_to_be_choosen)     # [0.54894334 0.4510567 ]
    word_choosen = np.random.choice(words_to_be_choosen, p=probs_to_be_choosen)
    print(bert_tokenizer.decode(word_choosen))
